
package org.apache.solr.handler.loader;

import java.io.IOException;
import java.io.Reader;
import java.io.StringReader;
import java.util.*;
import org.apache.commons.io.IOUtils;
import org.apache.noggit.JSONParser;
import org.apache.noggit.ObjectBuilder;
import org.apache.solr.common.SolrException;
import org.apache.solr.common.SolrInputDocument;
import org.apache.solr.common.SolrInputField;
import org.apache.solr.common.params.*;
import org.apache.solr.common.util.ContentStream;
import org.apache.solr.handler.RequestHandlerUtils;
import org.apache.solr.handler.UpdateRequestHandler;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.response.SolrQueryResponse;
import org.apache.solr.update.AddUpdateCommand;
import org.apache.solr.update.CommitUpdateCommand;
import org.apache.solr.update.DeleteUpdateCommand;
import org.apache.solr.update.RollbackUpdateCommand;
import org.apache.solr.update.processor.UpdateRequestProcessor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @since solr 4.0
 */
public class JsonLoader extends ContentStreamLoader {

    final static Logger log = LoggerFactory.getLogger(JsonLoader.class);

    @Override
    public String getDefaultWT() {
        return "json";
    }

    @Override
    public void load(SolrQueryRequest req, SolrQueryResponse rsp, ContentStream stream, UpdateRequestProcessor processor) throws Exception {
        new SingleThreadedJsonLoader(req, processor).load(req, rsp, stream, processor);
    }

    static class SingleThreadedJsonLoader extends ContentStreamLoader {

        protected final UpdateRequestProcessor processor;
        protected final SolrQueryRequest req;
        protected JSONParser parser;
        protected final int commitWithin;
        protected final boolean overwrite;

        public SingleThreadedJsonLoader(SolrQueryRequest req, UpdateRequestProcessor processor) {
            this.processor = processor;
            this.req = req;

            commitWithin = req.getParams().getInt(UpdateParams.COMMIT_WITHIN, -1);
            overwrite = req.getParams().getBool(UpdateParams.OVERWRITE, true);
        }

        @Override
        public void load(SolrQueryRequest req, SolrQueryResponse rsp, ContentStream stream, UpdateRequestProcessor processor) throws Exception {

            Reader reader = null;
            try {
                reader = stream.getReader();
                if(log.isTraceEnabled()) {
                    String body = IOUtils.toString(reader);
                    log.trace("body", body);
                    reader = new StringReader(body);
                }

                parser = new JSONParser(reader);
                this.processUpdate();
            }
            finally {
                IOUtils.closeQuietly(reader);
            }
        }

        @SuppressWarnings("fallthrough")
        void processUpdate() throws IOException {

            int ev = parser.nextEvent();
            while(ev != JSONParser.EOF) {

                switch(ev) {
                    case JSONParser.ARRAY_START:
                        handleAdds();
                        break;

                    case JSONParser.STRING:
                        if(parser.wasKey()) {
                            String v = parser.getString();
                            switch(v) {
                                case UpdateRequestHandler.ADD:
                                    int ev2 = parser.nextEvent();
                                    if(ev2 == JSONParser.OBJECT_START) {
                                        processor.processAdd(parseAdd());
                                    }
                                    else if(ev2 == JSONParser.ARRAY_START) {
                                        handleAdds();
                                    }
                                    else {
                                        assertEvent(ev2, JSONParser.OBJECT_START);
                                    }
                                    break;
                                case UpdateRequestHandler.COMMIT: {
                                    CommitUpdateCommand cmd = new CommitUpdateCommand(req, false);
                                    cmd.waitSearcher = true;
                                    parseCommitOptions(cmd);
                                    processor.processCommit(cmd);
                                    break;
                                }
                                case UpdateRequestHandler.OPTIMIZE: {
                                    CommitUpdateCommand cmd = new CommitUpdateCommand(req, true);
                                    cmd.waitSearcher = true;
                                    parseCommitOptions(cmd);
                                    processor.processCommit(cmd);
                                    break;
                                }
                                case UpdateRequestHandler.DELETE:
                                    handleDeleteCommand();
                                    break;
                                case UpdateRequestHandler.ROLLBACK:
                                    processor.processRollback(parseRollback());
                                    break;
                                default:
                                    throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Unknown command: " + v + " [" + parser.getPosition() + "]");
                            }
                            break;
                        }
                    // fall through

                    case JSONParser.LONG:
                    case JSONParser.NUMBER:
                    case JSONParser.BIGNUMBER:
                    case JSONParser.BOOLEAN:
                    case JSONParser.NULL:
                        log.info("can't have a value here! " + JSONParser.getEventString(ev) + " " + parser.getPosition());

                    case JSONParser.OBJECT_START:
                    case JSONParser.OBJECT_END:
                    case JSONParser.ARRAY_END:
                        break;

                    default:
                        log.info("Noggit UNKNOWN_EVENT_ID:" + ev);
                        break;
                }
                // read the next event
                ev = parser.nextEvent();
            }
        }

        //
        // "delete":"id"
        // "delete":["id1","id2"]
        // "delete":{"id":"foo"}
        // "delete":{"query":"myquery"}
        //
        void handleDeleteCommand() throws IOException {

            int ev = parser.nextEvent();
            switch(ev) {
                case JSONParser.ARRAY_START:
                    handleDeleteArray(ev);
                    break;
                case JSONParser.OBJECT_START:
                    handleDeleteMap(ev);
                    break;
                default:
                    handleSingleDelete(ev);
            }
        }

        // returns the string value for a primitive value, or null for the null value
        String getString(int ev) throws IOException {

            switch(ev) {
                case JSONParser.STRING:
                    return parser.getString();
                case JSONParser.BIGNUMBER:
                case JSONParser.NUMBER:
                case JSONParser.LONG:
                    return parser.getNumberChars().toString();
                case JSONParser.BOOLEAN:
                    return Boolean.toString(parser.getBoolean());
                case JSONParser.NULL:
                    return null;
                default:
                    throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
                            "Expected primitive JSON value but got: " + JSONParser.getEventString(ev) + " at [" + parser.getPosition() + "]");
            }
        }

        void handleSingleDelete(int ev) throws IOException {

            if(ev == JSONParser.OBJECT_START) {
                handleDeleteMap(ev);
            }
            else {
                DeleteUpdateCommand cmd = new DeleteUpdateCommand(req);
                cmd.commitWithin = commitWithin;
                String id = getString(ev);
                cmd.setId(id);
                processor.processDelete(cmd);
            }
        }

        void handleDeleteArray(int ev) throws IOException {

            assert ev == JSONParser.ARRAY_START;
            for(;;) {
                ev = parser.nextEvent();
                if(ev == JSONParser.ARRAY_END) {
                    return;
                }
                handleSingleDelete(ev);
            }
        }

        void handleDeleteMap(int ev) throws IOException {

            assert ev == JSONParser.OBJECT_START;

            DeleteUpdateCommand cmd = new DeleteUpdateCommand(req);
            cmd.commitWithin = commitWithin;

            while(true) {
                ev = parser.nextEvent();
                if(ev == JSONParser.STRING) {
                    String key = parser.getString();
                    if(parser.wasKey()) {
                        switch(key) {
                            case "id":
                                cmd.setId(getString(parser.nextEvent()));
                                break;
                            case "query":
                                cmd.setQuery(parser.getString());
                                break;
                            case "commitWithin":
                                cmd.commitWithin = (int)parser.getLong();
                                break;
                            case "_version_":
                                cmd.setVersion(parser.getLong());
                                break;
                            default:
                                throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Unknown key: " + key + " [" + parser.getPosition() + "]");
                        }
                    }
                    else {
                        throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "invalid string: " + key + " at [" + parser.getPosition() + "]");
                    }
                }
                else if(ev == JSONParser.OBJECT_END) {
                    if(cmd.getId() == null && cmd.getQuery() == null) {
                        throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Missing id or query for delete [" + parser.getPosition() + "]");
                    }

                    processor.processDelete(cmd);
                    return;
                }
                else {
                    throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Got: " + JSONParser.getEventString(ev) + " at [" + parser.getPosition() + "]");
                }
            }
        }

        RollbackUpdateCommand parseRollback() throws IOException {

            assertNextEvent(JSONParser.OBJECT_START);
            assertNextEvent(JSONParser.OBJECT_END);

            return new RollbackUpdateCommand(req);
        }

        void parseCommitOptions(CommitUpdateCommand cmd) throws IOException {

            assertNextEvent(JSONParser.OBJECT_START);
            final Map<String, Object> map = (Map)ObjectBuilder.getVal(parser);

            // SolrParams currently expects string values...
            SolrParams p = new SolrParams() {
                @Override
                public String get(String param) {
                    Object o = map.get(param);
                    return o == null ? null : o.toString();
                }

                @Override
                public String[] getParams(String param) {
                    return new String[]{get(param)};
                }

                @Override
                public Iterator<String> getParameterNamesIterator() {
                    return map.keySet().iterator();
                }
            };

            RequestHandlerUtils.validateCommitParams(p);
            p = SolrParams.wrapDefaults(p, req.getParams());   // default to the normal request params for commit options
            RequestHandlerUtils.updateCommit(cmd, p);
        }

        AddUpdateCommand parseAdd() throws IOException {

            AddUpdateCommand cmd = new AddUpdateCommand(req);
            cmd.commitWithin = commitWithin;
            cmd.overwrite = overwrite;

            float boost = 1.0f;

            while(true) {
                int ev = parser.nextEvent();
                if(ev == JSONParser.STRING) {
                    if(parser.wasKey()) {
                        String key = parser.getString();
                        switch(key) {
                            case "doc":
                                if(cmd.solrDoc != null) {
                                    throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "multiple docs in same add command");
                                }
                                ev = assertNextEvent(JSONParser.OBJECT_START);
                                cmd.solrDoc = parseDoc(ev);
                                break;
                            case UpdateRequestHandler.OVERWRITE:
                                cmd.overwrite = parser.getBoolean(); // reads next boolean
                                break;
                            case UpdateRequestHandler.COMMIT_WITHIN:
                                cmd.commitWithin = (int)parser.getLong();
                                break;
                            case "boost":
                                boost = Float.parseFloat(parser.getNumberChars().toString());
                                break;
                            default:
                                throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Unknown key: " + key + " [" + parser.getPosition() + "]");
                        }
                    }
                    else {
                        throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Should be a key at [" + parser.getPosition() + "]");
                    }
                }
                else if(ev == JSONParser.OBJECT_END) {
                    if(cmd.solrDoc == null) {
                        throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "missing solr document. " + parser.getPosition());
                    }
                    cmd.solrDoc.setDocumentBoost(boost);
                    return cmd;
                }
                else {
                    throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Got: " + JSONParser.getEventString(ev) + " at [" + parser.getPosition() + "]");
                }
            }
        }

        void handleAdds() throws IOException {

            while(true) {
                AddUpdateCommand cmd = new AddUpdateCommand(req);
                cmd.commitWithin = commitWithin;
                cmd.overwrite = overwrite;

                int ev = parser.nextEvent();
                if(ev == JSONParser.ARRAY_END) {
                    break;
                }

                assertEvent(ev, JSONParser.OBJECT_START);
                cmd.solrDoc = parseDoc(ev);
                processor.processAdd(cmd);
            }
        }

        int assertNextEvent(int expected) throws IOException {

            int got = parser.nextEvent();
            assertEvent(got, expected);

            return got;
        }

        void assertEvent(int ev, int expected) {

            if(ev != expected) {
                throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
                        "Expected: " + JSONParser.getEventString(expected)
                        + " but got " + JSONParser.getEventString(ev)
                        + " at [" + parser.getPosition() + "]");
            }
        }

        private SolrInputDocument parseDoc(int ev) throws IOException {

            assert ev == JSONParser.OBJECT_START;

            SolrInputDocument sdoc = new SolrInputDocument();
            for(;;) {
                SolrInputField sif = parseField();
                if(sif == null) {
                    return sdoc;
                }
                SolrInputField prev = sdoc.put(sif.getName(), sif);
                if(prev != null) {
                    // blech - repeated keys
                    sif.addValue(prev.getValue(), prev.getBoost());
                }
            }
        }

        private SolrInputField parseField() throws IOException {

            int ev = parser.nextEvent();
            if(ev == JSONParser.OBJECT_END) {
                return null;
            }

            String fieldName = parser.getString();
            SolrInputField sif = new SolrInputField(fieldName);
            parseFieldValue(sif);
            return sif;
        }

        private void parseFieldValue(SolrInputField sif) throws IOException {

            int ev = parser.nextEvent();
            if(ev == JSONParser.OBJECT_START) {
                parseExtendedFieldValue(sif, ev);
            }
            else {
                Object val = parseNormalFieldValue(ev);
                sif.setValue(val, 1.0f);
            }
        }

        private void parseExtendedFieldValue(SolrInputField sif, int ev) throws IOException {

            assert ev == JSONParser.OBJECT_START;

            float boost = 1.0f;
            Object normalFieldValue = null;
            Map<String, Object> extendedInfo = null;

            for(;;) {
                ev = parser.nextEvent();
                switch(ev) {
                    case JSONParser.STRING:
                        String label = parser.getString();
                        switch(label) {
                            case "boost":
                                ev = parser.nextEvent();
                                if(ev != JSONParser.NUMBER && ev != JSONParser.LONG && ev != JSONParser.BIGNUMBER) {
                                    throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "boost should have number! " + JSONParser.getEventString(ev));
                                }
                                boost = (float)parser.getDouble();
                                break;
                            case "value":
                                normalFieldValue = parseNormalFieldValue(parser.nextEvent());
                                break;
                            default:
                                // If we encounter other unknown map keys, then use a map
                                if(extendedInfo == null) {
                                    extendedInfo = new HashMap<String, Object>(2);
                                }
                                // for now, the only extended info will be field values
                                // we could either store this as an Object or a SolrInputField
                                Object val = parseNormalFieldValue(parser.nextEvent());
                                extendedInfo.put(label, val);
                                break;
                        }
                        break;

                    case JSONParser.OBJECT_END:
                        if(extendedInfo != null) {
                            if(normalFieldValue != null) {
                                extendedInfo.put("value", normalFieldValue);
                            }
                            sif.setValue(extendedInfo, boost);
                        }
                        else {
                            sif.setValue(normalFieldValue, boost);
                        }
                        return;

                    default:
                        throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Error parsing JSON extended field value. Unexpected " + JSONParser.getEventString(ev));
                }
            }
        }

        private Object parseNormalFieldValue(int ev) throws IOException {

            if(ev == JSONParser.ARRAY_START) {
                List<Object> val = parseArrayFieldValue(ev);
                return val;
            }
            else {
                Object val = parseSingleFieldValue(ev);
                return val;
            }
        }

        private Object parseSingleFieldValue(int ev) throws IOException {

            switch(ev) {
                case JSONParser.STRING:
                    return parser.getString();
                case JSONParser.LONG:
                case JSONParser.NUMBER:
                case JSONParser.BIGNUMBER:
                    return parser.getNumberChars().toString();
                case JSONParser.BOOLEAN:
                    return Boolean.toString(parser.getBoolean()); // for legacy reasons, single values s are expected to be strings
                case JSONParser.NULL:
                    parser.getNull();
                    return null;
                case JSONParser.ARRAY_START:
                    return parseArrayFieldValue(ev);
                default:
                    throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Error parsing JSON field value. Unexpected " + JSONParser.getEventString(ev));
            }
        }

        private List<Object> parseArrayFieldValue(int ev) throws IOException {

            assert ev == JSONParser.ARRAY_START;

            ArrayList lst = new ArrayList(2);
            for(;;) {
                ev = parser.nextEvent();
                if(ev == JSONParser.ARRAY_END) {
                    return lst;
                }
                Object val = parseSingleFieldValue(ev);
                lst.add(val);
            }
        }
    }
}